import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.naive_bayes import CategoricalNB

# 读取数据集
data = {
    '年龄': ['≤30', '≤30', '31~40', '>40', '>40', '>40', '31~40', '≤30', '≤30', '>40', '≤30', '31~40', '31~40', '>40', '≤30'],
    '收入': ['高', '高', '高', '中', '低', '低', '低', '中', '低', '中', '中', '中', '高', '中', '中'],
    '是否为学生': ['否', '否', '否', '否', '是', '是', '是', '否', '是', '是', '是', '否', '是', '否', '是'],
    '信誉': ['中', '优', '中', '中', '中', '优', '优', '中', '中', '中', '优', '优', '中', '优', '中'],
    '是否购买计算机': ['否', '否', '是', '是', '是', '否', '是', '否', '是', '是', '是', '是', '是', '否', '是']
}
df = pd.DataFrame(data)

# 数据预处理
# 使用 LabelEncoder 对定性数据进行编码
label_encoder = LabelEncoder()
# 对年龄进行编码，≤30 编码为 0，31~40 编码为 1，>40 编码为 2
df['年龄'] = label_encoder.fit_transform(df['年龄'])
print("年龄编码对应关系：≤30 -> 0，31~40 -> 1，>40 -> 2")
# 对收入进行编码，高编码为 0，中编码为 1，低编码为 2
df['收入'] = label_encoder.fit_transform(df['收入'])
print("收入编码对应关系：高 -> 0，中 -> 1，低 -> 2")
# 对是否为学生进行编码，否编码为 0，是编码为 1
df['是否为学生'] = label_encoder.fit_transform(df['是否为学生'])
print("是否为学生编码对应关系：否 -> 0，是 -> 1")
# 对信誉进行编码，中编码为 0，优编码为 1
df['信誉'] = label_encoder.fit_transform(df['信誉'])
print("信誉编码对应关系：中 -> 0，优 -> 1")
# 对是否购买计算机进行编码，否编码为 0，是编码为 1
df['是否购买计算机'] = label_encoder.fit_transform(df['是否购买计算机'])
print("是否购买计算机编码对应关系：否 -> 0，是 -> 1")

# 划分训练集和测试集
X = df.drop('是否购买计算机', axis=1)
y = df['是否购买计算机']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建朴素贝叶斯分类器
clf = CategoricalNB()

# 训练分类器
clf.fit(X_train, y_train)

# 新样本数据
X1 = [[2, 0, 1, 1]]  # 年龄：>40,收入：高，是否为学生：是，信誉：优
X2 = [[1, 2, 0, 0]]  # 年龄：31~40,收入：低，是否为学生：否，信誉：中

# 对新样本进行分类
prediction1 = clf.predict(X1)
prediction2 = clf.predict(X2)

print(f"新样本 X1 的分类结果：{prediction1[0]}")
print(f"新样本 X2 的分类结果：{prediction2[0]}")